import os
import datetime
import matplotlib.pyplot as plt
import numpy as np
import scipy.signal

class LossHistoryN:
    """
    可以记录任意数目的损失、准确率
    """
    def __init__(self, root, loss_name: [list, tuple] = ('train-loss', 'val-loss','train-acc','val-acc',), path=None):
        """
        :param root: 保存数据的根目录
        :param loss_name: 需要记录的损失的名字
        :param path: 本次保存的路径
        """
        # assert 是满足条件就不打印东西
        assert len(loss_name) != 0, "Loss name shouldn't be empty."
        self.root = self.__check_dir(root)
        self.loss_name = loss_name

        if path is None:
            current_time = datetime.datetime.now()
            str_time = datetime.datetime.strftime(current_time, '%Y_%m_%d_%H_%M_%S')
            self.str_time = str_time
        else:
            self.str_time = path

        self.dir_name = r'loss_%s' % self.str_time
        self.loss_dir = self.__check_dir(os.path.join(self.root, self.dir_name))

        # 将数据列表和保存数据的路径的信息保存到字典中
        self.dictionary = {}  # 保存数据
        self.paths = {}  # 保存数据的TXT路径
        for name in loss_name:
            self.dictionary[name] = []
            self.paths[name] = os.path.join(self.loss_dir, '%s.txt' % name)

        self.length = 0#数量
        self.colors = ['red', 'y', 'blue', 'k', 'green', 'm', 'black', 'c']
        self.linestyles = ['--', '-.', ':', '-']

    def add_loss(self, **kwargs):
        """
        :param kwargs: 一个参数字典，输入是需要字典的键值和初始化时self.loss_name中的一样才会正常运行
        :return:
        """
        assert len(kwargs.keys()) == len(self.loss_name), 'kwargs name must equal loss name.'
        names = kwargs.keys()
        # 添加数据到字典中
        for name in names:
            assert name in self.loss_name, 'kwargs.keys must in self.loss_name.'
            loss_value = kwargs[name]
            self.dictionary[name].append(loss_value)
            self.save_loss(self.paths[name], loss_value)

        self.length += 1
        # self.plt_loss()

    def plt_loss(self):
        length = range(self.length)

        a= np.max(self.dictionary[self.loss_name[0]])

        fig = plt.figure('loss')
        ax = fig.add_subplot(111)
        ax.plot(length, self.dictionary[self.loss_name[0]], '-', label='train_loss')
        ax.plot(length, self.dictionary[self.loss_name[1]], '-', label='val_loss')

        ax2 = ax.twinx()
        ax2.plot(length, self.dictionary[self.loss_name[2]], '-r', label='train_acc')
        ax2.plot(length, self.dictionary[self.loss_name[3]], '-r', label='val_acc')
        # ax.legend(loc=1)
        ax.legend(loc="upper left")

        ax.grid()
        ax.set_xlabel("Epoch")
        ax.set_ylabel(r"Loss")
        ax2.set_ylabel(r"ACC")
        # ax2.set_ylim(0, 100)
        # ax.set_ylim(0, 100)
        ax2.legend(loc=1)

        # 保存图片
        plt.savefig(os.path.join(self.loss_dir, "epoch_loss_%s.png" % self.str_time))
        plt.close('loss')

    @staticmethod
    def save_loss(save_path, loss):
        with open(save_path, 'a', encoding='utf-8') as file:
            file.write(str(loss) + '\n')

    @staticmethod
    def __check_dir(_path):
        if not os.path.exists(_path):
            os.makedirs(_path)
        return _path


if __name__ == '__main__':

    loss_history = LossHistoryN('./loss', loss_name=('train_loss', 'val_loss','train_acc','val_acc'))
    for i in [15, 45, 7, 12, 49, 66, 18, 45, 18, 65, 35, 57]:
        loss_history.add_loss(train_loss=i / 2, val_loss=i * 2 / 5,train_acc = i * 2 / 3,val_acc = 1)
    loss_history.plt_loss()

